#!/usr/bin/env python

############################################################
#1. This file loads the output file from pdb_parse.py
#2. generates a SDR table acroding to the set conditions
#3. calls the predictor to generate scores
#(conditions can be loop in this script)
############################################################
from Bio import AlignIO
#from Bio import SeqIO
#from Bio.PDB import *
#from Bio.SeqRecord import SeqRecord
#from Bio.Seq import Seq
#import os
#import re
import sys
#import MySQLdb
#import hmmer3
import pprint
import pickle
from collections import defaultdict
from bio import blosum
from itertools import combinations
import PDZ_predictor_valid
import csv
import sqlite3

def dbinsert(ins, dbname):
    print "Updating SDR into database:", dbname
    db = sqlite3.connect(dbname)
    curs = db.cursor()
    curs.execute("""delete from sdr""")
    for p in ins:
        print "Position", p, "SDR:", ins[p]
        for i in ins[p]:
            curs.execute("""insert into sdr
                         values (?, ?, ?)""", (i[0], i[1], i[2]))
    db.commit()

def read_inputfile(contactfile = 'contactinfo.dat', alnfile = 'tmp_aln.fasta'):
    #contactinfo.dat generate by ?
    with open(contactfile, 'rb') as pkl_file: 
        contactdat, pos_glgf, pos_gd = pickle.load(pkl_file)
    
        aln = AlignIO.read(alnfile, "fasta")
        #TODO: read contact information from database
        
        #pprint.pprint (contactdat)
        #raw_input()
        return contactdat, pos_glgf, pos_gd, aln

def is_similar(aln_column, similar_pt, submatrix = 'blosum62'):
    #aln_column = 1 column of sequence alignment
    #similar_pt = the test similar level
    #submatrix = sub matrix
    #It calculates how many residues in the column
    #will result in the same group of SDRs. If the
    #% is larger than similar level, return False.
    #else True
    #print aln_column
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    subsmat = blosum.Matrix(submatrix)
    res_count = defaultdict(int)
    
    for aa in aln_column:
        res_count[aa] += 1
    #print res_count
    for aa in res_count.keys():
        similar_count = 0
        if aa in AminoAcids and res_count[aa] > 0:
            similar_res_group = subsmat.substitutions(aa, 0)
            #similar_count = res_count[aa]
            for res in similar_res_group:
                similar_count += res_count[res]
            
            similar_level = float(similar_count) / float(len(aln_column))
            #print similar_count, len(aln_column)
            #print aa, similar_level
            if similar_level > similar_pt:
                #print "return true!"
                return True
    
    return False

def select_sdr_position(contact, pos_glgf, pos_gd, aln, condition_contact_pt,
               condition_similarres_pt, submatrix, dist, peppos):
    sdr_column = []
    total_alignments = len(aln)
    for motif in contact[peppos]:
        #print peppos, motif
        #print 'HERE:', contact[peppos][motif]
        sdr_count = 0
        for sdr in contact[peppos][motif]:
            #print type(sdr[2])
            if sdr[2]<= dist:
                sdr_count += 1
                
        if sdr_count > total_alignments * condition_contact_pt:
            #output format: peptide#, motif, offset
            
            if motif[0] == 'GLGF':
                #print 'GLGF', pos_glgf, motif[1]
                index = pos_glgf + motif[1] -1
            elif motif[0] == 'GD':
                #print 'GD', pos_gd, motif[1]
                index = pos_gd + motif[1] -1
            else:
                print 'Error: not GLGF or GD motif'
            if is_similar(aln.get_column(index), condition_similarres_pt, submatrix) is not True:
                #is_similar return True if the similarity is greater than
                #the condition
                #return ([motif[0], motif[1],peppos])
                sdr_column.append([motif[0], motif[1],peppos])
    
    if len(sdr_column) > 0:
        return sdr_column
    else:
        return False
        

def select_sdr(contact, pos_glgf, pos_gd, aln, condition_contact_pt,
               condition_similarres_pt, submatrix, dist):
    #1. total number of alignment
    #2. if column in aln similar not use
    #3. if len 'motif + offset' > total * thres use
    #4. import into mysql
    
    
    sdr_column = []
    #for i in range(aln.get_alignment_length()):
    #    aln.get_column(i)
    for peppos in contact:
        position_sdr = select_sdr_position(contact, pos_glgf, pos_gd, aln, condition_contact_pt,
                                           condition_similarres_pt, submatrix, dist, peppos)
        print position_sdr
        if position_sdr is not False:
            sdr_column += position_sdr
        else:
            return False
        
    return sdr_column            
    
def sdr_check(sdrs, lenpos = 5):
    #check if there is at least an SDR for every position
    position = set()
    for sdr in sdrs:
        #print sdr[2]
        position.add(sdr[2])
    #print len(position)
    if len(position) == lenpos:
        return True
    else:
        return False


def select_sdr2(cond1_range, cond2_range, subM_range, dist_range, dbname):
    #cond1 = min% of contacting residues in a column (> cond1 will return true)
    #cond2 = max% of similar residues in a column (< cond2 will return true)
    #subM = substitution matrix use to calculate similar residues
    #print "Reading contact file"
    total_sdr = {}
    contactdata, pos_glgf, pos_gd, aln = read_inputfile()
    for peppos in contactdata:
        total_sdr[peppos] = dict()
        for cond1 in cond1_range:
            for cond2 in cond2_range:
                for subM in subM_range:
                    for dist in dist_range:
                        #
                        key = (cond1, cond2, subM, dist)
                        
                        sdr_row = select_sdr_position(contactdata, pos_glgf, pos_gd, aln, cond1,
                                           cond2, subM, dist, peppos)
                        
                        #print sdr_row
                        if sdr_row is not False:
                            print peppos, key
                            total_sdr[peppos][key] = sdr_row
    return total_sdr

def select_sdr3(cond1, cond2, subM, dist, dbname):
    total_sdr = select_sdr2([cond1], [cond2], [subM], [dist], dbname)
    
    key = (cond1, cond2, subM, dist)
    output = dict()
    is_valid = True
    for peppos in total_sdr:
        try:
            output[peppos] = total_sdr[peppos][key]
        except KeyError:
            output[peppos] = []
            is_valid = False
    if is_valid is True:
        return output
    else:
        #raise
        return False
        
        
        
'''
def select_sdr3(total_sdr):
    for key in total_sdr:
        
    #output_sdr = select_sdr(contactdata, pos_glgf, pos_gd,
    #                        aln, cond1, cond2, subM, dist)
    
    print 'Length of Output SDR:', len(output_sdr)
    if len(output_sdr) < 1 or not (sdr_check(output_sdr)):
        print "Not engouth SDR"
        print output_sdr
        return False, False
    
    elif output_sdr != []:
        print "Selected SDRs are:"
        pprint.pprint(output_sdr)
        dbinsert(output_sdr, dbname)
        
        return True
        #roc, auc = PDZ_predictor_valid.main(validtable, dbname)
        #return roc, auc
        
    else:
        print "No SDR selected."
        return False, False
'''   
def validate():
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    #contact_percent = 0.3
    #similar_percent = 0.6
    #subm = 'blosum62'
    #dist = 6
    submrange = ['blosum62'] #substitution matrix
    distrange = [10] #contact distance
    cptrange = [0.4] #min contact res % in SDR column
    sptrange = [0.6] #max similar res % in SDR column
    for dist in distrange:
        for contact_percent in cptrange:
            for similar_percent in sptrange:
                for subm in submrange:
                    for dbnum in range(0, 10):  #debug
                        print "#------------------------------------#"
                        print "#Testing parameters#"
                        print "Substitution Matrix:",subm
                        print "Max. Similar Res%:", similar_percent
                        print "Min. Contact Res%:", contact_percent
                        print "Max. Contact Distance:", dist
                        print "On validation database:", dbnum + 1
                        
                        dbname = 'pdzvalid_' + str(dbnum + 1)
                        validtable = 'test_group'+str(dbnum)+'.dat'
                        #a, b = main(validtable, dbname);
                        rocoord, auc = main(contact_percent, similar_percent,
                                            subm, dist, dbname, validtable)
                        if rocoord:
                            fname = ('ROC_'+str(dist)+'_'+str(int(contact_percent*100))+
                                     '_'+str(int(similar_percent*100))+'_'+subm+'_'+
                                     str(dbnum)+'.output')
                            with open(fname, 'w') as f:
                                for i in rocoord:
                                    writer = csv.writer(f, delimiter = '\t',
                                                        lineterminator = '\n')
                                    writer.writerow(i)
                            with open('AUC.out', 'a') as f:
                                writer = csv.writer(f, delimiter = '\t',
                                                    lineterminator = '\n')
                                writer.writerow([dist, contact_percent, similar_percent,
                                                 subm, dbnum, auc])
                        else:
                            print "Return false from predictor"
                            
                            #print "predictor return null"
    #print rocoord
    #print 'AUC:' auc
    #PDZ_predictor_valid.main()
    
def normal():
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    #contact_percent = 0.3
    #similar_percent = 0.6
    #subm = 'blosum62'
    #dist = 6
    submrange = ['blosum62', 'blosum80'] #substitution matrix
    distrange = [6,8,10] #contact distance
    cptrange = [0.4,0.6,0.8] #min contact res % in SDR column
    sptrange = [0.6,0.6,0.8] #max similar res % in SDR column
    dbname = "pdzvalid.db"
    contact_percent = 0.6
    similar_percent = 0.8
    subm = 'blosum62'
    dist = 8
    print "#------------------------------------#"
    print "#Testing parameters#"
    print "Substitution Matrix:",subm
    print "Max. Similar Res%:", similar_percent
    print "Min. Contact Res%:", contact_percent
    print "Max. Contact Distance:", dist

    #conditions:
    #contact percentage, similarity percentager, sub matrix, distance 
    
    flag_sdr = select_sdr3(contact_percent, similar_percent,
                subm, dist, dbname)
    
    '''
    if flag_sdr is True:
        peptides = ["VFKLT","KVVLK","KTGTL"]
        seq = PDZ_predictor_valid.fasta2seqrecord("test.fasta")
        scores, pdzcount = PDZ_predictor_valid.scoring(seq, peptides, dbname)
    else:
        print "-----------SDR ERROR"
    print scores, pdzcount
    '''

def load_negative(infile):
    #convert negative binding data generated by nonbind_random.py to
    #a usable format?
    return

def score_nagative():


    return 

if __name__ == '__main__':
    
    normal()